package com.gitee.c0dehvb.reactive.sample.r2dbc.connection;

import com.gitee.c0dehvb.reactive.sample.r2dbc.ShardingRewriteEngine;
import io.r2dbc.spi.Connection;
import io.r2dbc.spi.ConnectionFactory;
import io.r2dbc.spi.ConnectionFactoryMetadata;
import org.reactivestreams.Publisher;
import reactor.core.publisher.Mono;

/**
 * 分片连接工厂
 * @author LiYilin
 * @since 2022/6/15 11:59 AM
 **/
public class ShardingConnectionFactory implements ConnectionFactory {
    private final ConnectionFactory connectionFactory;
    private final ShardingRewriteEngine shardingRewriteEngine;

    public ShardingConnectionFactory(ConnectionFactory connectionFactory, ShardingRewriteEngine shardingRewriteEngine) {
        this.connectionFactory = connectionFactory;
        this.shardingRewriteEngine = shardingRewriteEngine;
    }

    @Override
    public Publisher<? extends Connection> create() {
        return Mono.from(this.connectionFactory.create())
                // 包装成ShardingConnection
                .map(connection -> new ShardingConnection(connection, shardingRewriteEngine));
    }

    @Override
    public ConnectionFactoryMetadata getMetadata() {
        return connectionFactory.getMetadata();
    }
}
